#include <assert.h>
#include <elf.h>
#include <errno.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define IMAGE_FILE "./image"
#define ARGS "[--extended] [--vm] <bootblock> <executable-file> ..."

#define Elf32_EHDR_SIZE sizeof(Elf32_Ehdr)
#define Elf32_PHDR_SIZE sizeof(Elf32_Phdr)
#define OS_SIZE_LOC 0x1fe


/* structure to store command line options */
static struct {
    int vm;
    int extended;
} options;

/* prototypes of local functions */
static void create_image(int nfiles, char *files[]);
static void error(char *fmt, ...);
static void read_ehdr(Elf32_Ehdr * ehdr, FILE * fp);
static void read_phdr(Elf32_Phdr * phdr, FILE * fp, int ph,
                      Elf32_Ehdr ehdr);
static void write_segment(Elf32_Ehdr ehdr, Elf32_Phdr phdr, FILE * fp,
                          FILE * img, int *nbytes, int *first);
static void write_os_size(int first, int pre_nbytes, int nbytes, FILE * img);

int main(int argc, char **argv)
{
    char *progname = argv[0];

    /* process command line options */
    options.vm = 0;
    options.extended = 0;
    while ((argc > 1) && (argv[1][0] == '-') && (argv[1][1] == '-')) {
        char *option = &argv[1][2];

        if (strcmp(option, "vm") == 0) {
            options.vm = 1;
        } else if (strcmp(option, "extended") == 0) {
            options.extended = 1;
        } else {
            error("%s: invalid option\nusage: %s %s\n", progname,
                  progname, ARGS);
        }
        argc--;
        argv++;
    }
    if (options.vm == 1) {
        error("%s: option --vm not implemented\n", progname);
    }
    if (argc < 2) {
        /* at least 3 args (createimage bootblock kernel) */
        error("usage: %s %s\n", progname, ARGS);
    }
    create_image(argc - 1, argv + 1);
    return 0;
}

static void create_image(int nfiles, char *files[])
{
    int ph, nbytes = 512, first = 0, pre_nbytes, fill_zero;
    FILE *fp, *img;
    Elf32_Ehdr ehdr;
    Elf32_Phdr phdr;

    /* open the image file */
    img = fopen(IMAGE_FILE, "wb");
    if (img == NULL) {
        printf("Create image ERROR!\n");
        return;
    }
    // fseek(img, 512, SEEK_SET);

    /* for each input file */
    while (nfiles-- > 0) {
        pre_nbytes = nbytes;

        /* open input file */
        fp = fopen(*files, "rb");
        if (fp == NULL) {
            printf("Read file %s ERROR!\n", *files);
            return;
        }
        /* read ELF header */
        read_ehdr(&ehdr, fp);
        printf("0x%04x: %s\n", ehdr.e_entry, *files);

        /* for each program header */
        for (ph = 0; ph < ehdr.e_phnum; ph++) {

            /* read program header */
            read_phdr(&phdr, fp, ph, ehdr);

            /* print some information */
            if (options.extended) {
                printf("\tsegment %d:\n", ph + 1);
                printf("\t\toffset:\t0x%x\tvaddr:\t0x%x\n", phdr.p_offset, phdr.p_vaddr);
                printf("\t\tfilezs:\t0x%x\tmemsz:\t0x%x\n", phdr.p_filesz, phdr.p_memsz);
            }

            /* write segment to the image */
            write_segment(ehdr, phdr, fp, img, &nbytes, &first);
        }

        /* fill the sector with 0 to align */
        fill_zero = 512 - (nbytes % 512);
        nbytes += fill_zero;
        if (fill_zero != 0) {
            fseek(img, nbytes, SEEK_SET);
            while (fill_zero > 0) {
                fputc(0, img);
                fill_zero--;
            }
        }

        first++;
        write_os_size(first, pre_nbytes, nbytes, img);
        
        fclose(fp);
        files++;
    }
    
    fclose(img);
}

inline static void read_ehdr(Elf32_Ehdr * ehdr, FILE * fp)
{
    rewind(fp);
    if (!fread(ehdr, Elf32_EHDR_SIZE, 1, fp)) {
        printf("Read ehdr ERROR!\n");
    }
    return;
}

inline static void read_phdr(Elf32_Phdr * phdr, FILE * fp, int ph,
                      Elf32_Ehdr ehdr)
{
    fseek(fp, ehdr.e_phoff + Elf32_PHDR_SIZE * ph, SEEK_SET);
    if (!fread(phdr, Elf32_PHDR_SIZE, 1, fp)) {
        printf("Read phdr ERROR!\n");
    }
    return;
}

static void write_segment(Elf32_Ehdr ehdr, Elf32_Phdr phdr, FILE * fp,
                          FILE * img, int *nbytes, int *first)
{
    int32_t ch, i;
    uint32_t filesz, memsz, offset;

    filesz = phdr.p_filesz;
    memsz = phdr.p_memsz;
    offset = phdr.p_offset;

    if (memsz == 0) {
        return;
    }

    fseek(fp, offset, SEEK_SET);
    fseek(img, *nbytes, SEEK_SET);

    for (i = 0; i < filesz; i++) {
        ch = fgetc(fp);
        fputc(ch, img);
    }

    for (; i < memsz; i++) {
        fputc(0, img);
    }

    /* renew first and nbytes */
    *nbytes += memsz;
    
    return;
}

static void write_os_size(int first, int pre_nbytes, int nbytes, FILE * img)
{
    short sector;
    int kernel_size;

    fseek(img, OS_SIZE_LOC, SEEK_SET);
    
    /* kernel_size is the sector number of the kernel */
    kernel_size = (nbytes - pre_nbytes) >> 9;

    /* tranform int to char */
    sector = (short)kernel_size;
    
    fwrite(&sector, 2, 1, img);
    if (options.extended) {
        printf("\tkernel size: %d sector(s)\n", kernel_size);
    }
    return;
}

/* print an error message and exit */
static void error(char *fmt, ...)
{
    va_list args;

    va_start(args, fmt);
    vfprintf(stderr, fmt, args);
    va_end(args);
    if (errno != 0) {
        perror(NULL);
    }
    exit(EXIT_FAILURE);
}
